from dotenv import load_dotenv
import os
load_dotenv()
OPENAI_BASEURL = os.environ.get('OPENAI_BASEURL', 'https://api.openai.com/v1')
OPENAI_SK = os.environ.get('OPENAI_SK', '')

language_model = os.environ.get('LANGUAGE_MODEL','gpt-3.5-turbo-1106')
INPUT_MAX_LENGTH = int(os.environ.get('INPUT_MAX_LENGTH','8000'))

if not OPENAI_SK:
  raise Exception('OPENAI_SK is not set')

import argparse
argparser = argparse.ArgumentParser(description='GPT Translater')

argparser.add_argument('--input-dir', type=str, required=True, help='输入数据文件夹')
argparser.add_argument('--input-file-pattern', type=str, default='**/*.tex', help='查找文件的pattern')
argparser.add_argument('--output-dir', type=str, required=False,help='输出数据文件夹')

args = argparser.parse_args()
args.input_dir = os.path.abspath(args.input_dir)
if not args.output_dir:
  input_dir_name = os.path.basename(args.input_dir)
  args.output_dir = os.path.abspath(os.path.join(args.input_dir, '..', f'{input_dir_name}_output'))

import glob
# 查找所有需要处理的文件
filename_list=[]
if os.path.exists(args.input_dir):
  print(f'要处理的文件列表 {args.input_dir}/{args.input_file_pattern}: ')  
  files = glob.glob(args.input_file_pattern, root_dir=args.input_dir, recursive=True)
  for f in files:
    print(f)
    filename_list.append(f)
else:
  raise Exception(f'输入文件夹不存在: {args.input_dir}')

finished_filename_list=[]
if os.path.exists(args.output_dir):
  print(f'已经处理过的文件列表 {args.output_dir}/{args.input_file_pattern}: ')  
  files = glob.glob(args.input_file_pattern, root_dir=args.output_dir, recursive=True)
  for f in files:
    print(f)
    finished_filename_list.append(f)


# yes = input('是否继续...(Y/n)')
# if yes.strip() and yes.strip().lower() != 'y':
#   exit(0)

from openai import OpenAI
client = OpenAI(api_key = OPENAI_SK, base_url=OPENAI_BASEURL)


def translate(modelname: str, latex_text: str, try_times=5):
  while try_times > 0:
    try:
      response = client.chat.completions.create(
        model=modelname,
        messages=[
          {"role": "system", "content": "请把LaTeX中的英文翻译为中文，注意保留LaTeX格式。比如label, citet, citep,  lstlisting, verbatim、数学公式里面内容的不要翻译。",},
          {"role": "user", "content": latex_text,},
        ],
      )
      choice = response.choices[0]
      text = choice.message.content
      if not text:
        print(response)
        raise Exception('返回为空')
      else:
        return text
    except Exception as e:
      print(f'翻译失败，重试中...{try_times}次机会')
      print(e)
      try_times -= 1
  raise Exception('翻译失败')

def chop_text(full_text: str, splitter='\n\n', max_length: int=INPUT_MAX_LENGTH):
  text_list = full_text.split(splitter)
  text = '' # 拼接好的文本
  for t in text_list:
    try_extend_text = (text + splitter + t).strip() # 尝试拼接文本，减少调用次数
    if len(try_extend_text) < max_length: # 说明还可以继续拼接
      text = try_extend_text
      continue
    elif not text: # corner case: 如果一开始就超过了 short_length。会导致yield为空串，因此避免空串
      text = t  
      continue
    else:
      # 此时 try_extend_text 拼接过长了，我们返回上一步的 text
      yield language_model, text
      # 本轮的 t 还没有处理，放入到text中
      text = t

  if not text:
    return

  yield language_model, text

from pydantic import BaseModel
import yaml
class TranslateStatus(BaseModel):
  current_in_file: str = ''
  current_out_file: str = ''
  current_idx: int = 0
  translated_blocks: list[str] = []

global_status = TranslateStatus()
STATUS_FILE = 'status.yaml'
if os.path.exists(STATUS_FILE):
  with open(STATUS_FILE, 'r', encoding='utf-8') as f:
    global_status = TranslateStatus(**yaml.load(f, Loader=yaml.SafeLoader))

for fname in filename_list:
  if fname in finished_filename_list:
    print(f'文件 {fname} 已经处理过了，跳过...')
    continue

  if global_status.current_in_file and global_status.current_in_file != os.path.join(args.input_dir, fname):
    print(f'缓存点： {global_status.current_in_file} 还没有处理完，跳过当前文件...')
    continue
  elif not global_status.current_in_file:
    global_status.current_in_file = os.path.join(args.input_dir, fname)
    global_status.current_out_file = os.path.join(args.output_dir, fname)

  assert(global_status.current_in_file==os.path.join(args.input_dir, fname))

  with open(global_status.current_in_file, 'r') as f:
    print(f'正在翻译文件: {fname}...')
    latex_text = f.read()
    results = global_status.translated_blocks
    for idx, (modelname, t) in enumerate(chop_text(latex_text)):
      if idx < global_status.current_idx:
        continue
      global_status.current_idx = idx

      t = t.strip()
      if not t:
        continue

      try:
        print('-'*20)
        print(f'翻译前: \n{t}')
        transed = translate(modelname, t)
        print(f'翻译后: \n{transed}')
        results.append(transed)
      except Exception as e:
        with open(STATUS_FILE, 'w', encoding='utf-8') as f:
          yaml.dump(global_status.model_dump(), f, allow_unicode=True)
        print(f'翻译失败，退出程序。状态缓存为 {STATUS_FILE}')
        exit(1)

    if not os.path.exists(args.output_dir):
      os.makedirs(args.output_dir)
      
    outfname = global_status.current_out_file
    if not os.path.exists(os.path.dirname(outfname)):
      os.makedirs(os.path.dirname(outfname))
    with open(outfname, 'w', encoding='utf-8') as f:
      f.write('\n\n'.join(results))

    global_status.current_idx = 0
    global_status.current_in_file = ''
    global_status.current_out_file = ''
    global_status.translated_blocks = []
